"""
时序数据模型
============

定义市场数据的时序存储模型，用于InfluxDB存储
"""

from datetime import datetime
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass, field
from enum import Enum
import json


class DataFrequency(str, Enum):
    """数据频率枚举"""
    TICK = "tick"           # 逐笔数据
    SECOND = "1s"           # 秒级
    MINUTE = "1m"           # 分钟级
    FIVE_MINUTE = "5m"      # 5分钟
    FIFTEEN_MINUTE = "15m"  # 15分钟
    THIRTY_MINUTE = "30m"   # 30分钟
    HOUR = "1h"             # 小时级
    FOUR_HOUR = "4h"        # 4小时
    DAILY = "1d"            # 日级
    WEEKLY = "1w"           # 周级
    MONTHLY = "1M"          # 月级
    
    @classmethod
    def from_string(cls, frequency_str: str) -> 'DataFrequency':
        """从字符串创建频率枚举"""
        for freq in cls:
            if freq.value == frequency_str:
                return freq
        raise ValueError(f"不支持的频率: {frequency_str}")


@dataclass
class TimeSeriesPoint:
    """时序数据点基类"""
    timestamp: datetime
    symbol: str
    measurement: str = ""
    tags: Dict[str, str] = field(default_factory=dict)
    fields: Dict[str, Union[float, int, str, bool]] = field(default_factory=dict)
    
    def to_influx_point(self) -> Dict[str, Any]:
        """转换为InfluxDB数据点格式"""
        return {
            'measurement': self.measurement,
            'time': self.timestamp,
            'tags': {
                'symbol': self.symbol,
                **self.tags
            },
            **self.fields
        }
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'timestamp': self.timestamp.isoformat(),
            'symbol': self.symbol,
            'measurement': self.measurement,
            'tags': self.tags,
            'fields': self.fields
        }


@dataclass
class TickData(TimeSeriesPoint):
    """Tick数据模型"""
    price: float = 0.0
    volume: float = 0.0
    bid_price: Optional[float] = None
    ask_price: Optional[float] = None
    bid_volume: Optional[float] = None
    ask_volume: Optional[float] = None
    
    def __post_init__(self):
        self.measurement = "tick_data"
        self.fields = {
            'price': self.price,
            'volume': self.volume,
            'bid_price': self.bid_price,
            'ask_price': self.ask_price,
            'bid_volume': self.bid_volume,
            'ask_volume': self.ask_volume
        }
        # 移除None值
        self.fields = {k: v for k, v in self.fields.items() if v is not None}


@dataclass
class KLineData(TimeSeriesPoint):
    """K线数据模型"""
    open_price: float = 0.0
    high_price: float = 0.0
    low_price: float = 0.0
    close_price: float = 0.0
    volume: float = 0.0
    turnover: Optional[float] = None
    frequency: DataFrequency = DataFrequency.MINUTE
    
    def __post_init__(self):
        self.measurement = "kline_data"
        self.tags['frequency'] = self.frequency.value
        self.fields = {
            'open': self.open_price,
            'high': self.high_price,
            'low': self.low_price,
            'close': self.close_price,
            'volume': self.volume,
            'turnover': self.turnover
        }
        # 移除None值
        self.fields = {k: v for k, v in self.fields.items() if v is not None}


@dataclass
class OrderBookData(TimeSeriesPoint):
    """订单簿数据模型"""
    bids: List[List[float]] = field(default_factory=list)  # [[price, volume], ...]
    asks: List[List[float]] = field(default_factory=list)  # [[price, volume], ...]
    depth: int = 20
    
    def __post_init__(self):
        self.measurement = "orderbook_data"
        self.tags['depth'] = str(self.depth)
        
        # 将订单簿数据序列化为JSON字符串存储
        self.fields = {
            'bids_json': json.dumps(self.bids),
            'asks_json': json.dumps(self.asks),
            'best_bid': self.bids[0][0] if self.bids else None,
            'best_ask': self.asks[0][0] if self.asks else None,
            'bid_volume': sum(bid[1] for bid in self.bids),
            'ask_volume': sum(ask[1] for ask in self.asks)
        }
        # 移除None值
        self.fields = {k: v for k, v in self.fields.items() if v is not None}


@dataclass
class TradeData(TimeSeriesPoint):
    """成交数据模型"""
    price: float = 0.0
    volume: float = 0.0
    direction: str = "buy"  # "buy" or "sell"
    trade_id: Optional[str] = None
    
    def __post_init__(self):
        self.measurement = "trade_data"
        self.tags['direction'] = self.direction
        self.fields = {
            'price': self.price,
            'volume': self.volume,
            'trade_id': self.trade_id
        }
        # 移除None值
        self.fields = {k: v for k, v in self.fields.items() if v is not None}


@dataclass
class MarketSummaryData(TimeSeriesPoint):
    """市场汇总数据模型"""
    total_volume: float = 0.0
    total_turnover: float = 0.0
    price_change: float = 0.0
    price_change_percent: float = 0.0
    high_24h: float = 0.0
    low_24h: float = 0.0
    open_24h: float = 0.0
    close_price: float = 0.0
    
    def __post_init__(self):
        self.measurement = "market_summary"
        self.fields = {
            'total_volume': self.total_volume,
            'total_turnover': self.total_turnover,
            'price_change': self.price_change,
            'price_change_percent': self.price_change_percent,
            'high_24h': self.high_24h,
            'low_24h': self.low_24h,
            'open_24h': self.open_24h,
            'close_price': self.close_price
        }


class TimeSeriesStorage:
    """时序数据存储管理器"""
    
    def __init__(self, influxdb_manager):
        self.influxdb_manager = influxdb_manager
    
    async def store_tick_data(self, tick_data: TickData) -> bool:
        """存储Tick数据"""
        points = [tick_data.to_influx_point()]
        return await self.influxdb_manager.write_points(points)
    
    async def store_kline_data(self, kline_data: KLineData) -> bool:
        """存储K线数据"""
        points = [kline_data.to_influx_point()]
        return await self.influxdb_manager.write_points(points)
    
    async def store_orderbook_data(self, orderbook_data: OrderBookData) -> bool:
        """存储订单簿数据"""
        points = [orderbook_data.to_influx_point()]
        return await self.influxdb_manager.write_points(points)
    
    async def store_trade_data(self, trade_data: TradeData) -> bool:
        """存储成交数据"""
        points = [trade_data.to_influx_point()]
        return await self.influxdb_manager.write_points(points)
    
    async def store_market_summary(self, summary_data: MarketSummaryData) -> bool:
        """存储市场汇总数据"""
        points = [summary_data.to_influx_point()]
        return await self.influxdb_manager.write_points(points)
    
    async def batch_store(self, data_points: List[TimeSeriesPoint]) -> bool:
        """批量存储数据点"""
        points = [point.to_influx_point() for point in data_points]
        return await self.influxdb_manager.write_points(points)
    
    async def query_kline_data(self, symbol: str, frequency: DataFrequency,
                              start: datetime, end: Optional[datetime] = None,
                              limit: int = 1000) -> List[Dict[str, Any]]:
        """查询K线数据"""
        if not end:
            end = datetime.now()
        
        query = f'''
        from(bucket: "{self.influxdb_manager.bucket}")
        |> range(start: {start.isoformat()}Z, stop: {end.isoformat()}Z)
        |> filter(fn: (r) => r._measurement == "kline_data")
        |> filter(fn: (r) => r.symbol == "{symbol}")
        |> filter(fn: (r) => r.frequency == "{frequency.value}")
        |> limit(n: {limit})
        |> sort(columns: ["_time"])
        '''
        
        return await self.influxdb_manager.execute_query(query)
    
    async def query_tick_data(self, symbol: str, start: datetime,
                             end: Optional[datetime] = None,
                             limit: int = 10000) -> List[Dict[str, Any]]:
        """查询Tick数据"""
        if not end:
            end = datetime.now()
        
        query = f'''
        from(bucket: "{self.influxdb_manager.bucket}")
        |> range(start: {start.isoformat()}Z, stop: {end.isoformat()}Z)
        |> filter(fn: (r) => r._measurement == "tick_data")
        |> filter(fn: (r) => r.symbol == "{symbol}")
        |> limit(n: {limit})
        |> sort(columns: ["_time"])
        '''
        
        return await self.influxdb_manager.execute_query(query)
    
    async def get_latest_price(self, symbol: str) -> Optional[float]:
        """获取最新价格"""
        query = f'''
        from(bucket: "{self.influxdb_manager.bucket}")
        |> range(start: -1h)
        |> filter(fn: (r) => r._measurement == "tick_data" or r._measurement == "kline_data")
        |> filter(fn: (r) => r.symbol == "{symbol}")
        |> filter(fn: (r) => r._field == "price" or r._field == "close")
        |> last()
        '''
        
        results = await self.influxdb_manager.execute_query(query)
        if results:
            return results[0].get('value')
        return None
    
    async def calculate_price_change_24h(self, symbol: str) -> Dict[str, float]:
        """计算24小时价格变化"""
        query = f'''
        from(bucket: "{self.influxdb_manager.bucket}")
        |> range(start: -24h)
        |> filter(fn: (r) => r._measurement == "kline_data")
        |> filter(fn: (r) => r.symbol == "{symbol}")
        |> filter(fn: (r) => r.frequency == "1m")
        |> filter(fn: (r) => r._field == "close")
        |> first()
        '''
        
        start_results = await self.influxdb_manager.execute_query(query)
        
        current_price = await self.get_latest_price(symbol)
        
        if start_results and current_price:
            start_price = start_results[0].get('value')
            if start_price:
                change = current_price - start_price
                change_percent = (change / start_price) * 100
                return {
                    'price_change': change,
                    'price_change_percent': change_percent,
                    'current_price': current_price,
                    'start_price': start_price
                }
        
        return {
            'price_change': 0.0,
            'price_change_percent': 0.0,
            'current_price': current_price or 0.0,
            'start_price': 0.0
        }
